import streamlit as st
import os
from langchain_community.llms import Tongyi
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from streamlit_extras.switch_page_button import switch_page

# App title
st.set_page_config(page_title="Chatbot")

def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
    st.session_state.history = []

# Tongyi Credentials
with st.sidebar:
    st.title('Welcome To ChatBot')
    if st.button('Clear History'):
        clear_chat_history()
    # tongyi_api_key = os.getenv("OPENAI_API_KEY")
    tongyi_api_key = 'sk-24686fb38c3648eb9dbc356c757c2678'
    # st.subheader('Models and parameters')
    # selected_model = st.selectbox('Choose a model', ['qwen-plus'], key='selected_model')
    # temperature = st.slider('temperature', min_value=0.01, max_value=1.0, value=0.1, step=0.01)
    # top_p = st.slider('top_p', min_value=0.01, max_value=1.0, value=0.9, step=0.01)
    # max_length = st.slider('max_length', min_value=32, max_value=512, value=120, step=8)
    temperature = 0.1

    # Initialize model based on selection
    llm = Tongyi(
        model="qwen-plus",
        dashscope_api_key=tongyi_api_key,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 填写DashScope服务endpoint
        temperature=temperature
    )

# Store LLM generated responses
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]

# 初始化或加载历史记录
if "history" not in st.session_state:
    st.session_state.history = []

# Display or clear chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# 主体内容区域
# 这里可以是你的应用逻辑，比如用户交互产生的新记录会追加到历史记录中
def add_to_history(item):
    st.session_state.history.append(item)
    # print('st.session_state.history====', st.session_state.history)
    # # 清空容器并仅渲染最新的历史记录列表
    with st.sidebar:
            history_container = st.empty() if 'history_container' not in st.session_state else st.session_state.history_container
            # history_container.empty()
            for record in st.session_state.history:
                print('record====', record)
                history_container.text(record)
    # print('history_container',history_container)
   
def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
    # st.sidebar.button('Clear Chat History', on_click=clear_chat_history)

# Function for generating Tongyi response
def generate_tongyi_response(prompt_input):
    string_dialogue = "You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'."
    for dict_message in st.session_state.messages:
        if dict_message["role"] == "user":
            string_dialogue += "User: " + dict_message["content"] + "\n\n"
        else:
            string_dialogue += "Assistant: " + dict_message["content"] + "\n\n"

    prompts = PromptTemplate(template=string_dialogue, input_variables=["用中文"])

    output_parser = StrOutputParser()

    chain = prompts | llm | output_parser
    return chain

# User-provided prompt
if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)
        add_to_history(prompt)

# Generate a new response if last message is not from assistant
if st.session_state.messages[-1]["role"] != "assistant":
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            prompts = PromptTemplate(template=prompt, input_variables=["用中文"])

            output_parser = StrOutputParser()

            chain = prompts | llm | output_parser
            try:
                response = chain.invoke(input={"text": prompt})
                print(response)  # 或者处理 response 的方式，具体取决于它的结构
                placeholder = st.empty()
                placeholder.markdown(response)
            except Exception as e:
                print(f"An error occurred: {e}")
        message = {"role": "assistant", "content": response}
        st.session_state.messages.append(message)
        add_to_history(message['content'])
        print('st.session_state',st.session_state)
